# -*- coding: utf-8 -*-

import pickle
import numpy as np
import pandas as pd
import random
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras.models import load_model
from mpl_toolkits.mplot3d import Axes3D

def bin2dec(s):
    try:
        dec = int(s,2)
    # 发现 Python 3.8 版本的一个 bug。一个不明所以的错误。
    # Python 将我的字符串按字符一个一个拆开了，但是检查了差不多 2 小时呀/(ㄒoㄒ)/~~
    except:
        s = ''.join(s)
        dec = int(s,2)
    # l 为 2 进制位数。
    l = len(s)     
    num = (1+2*dec)/(2**(l+1)-1)
    return num

def rand_binary_vec(num_in_pop):
    '''
    随机产生一个向量，向量的长度为 num_in_pop(种群个体数)
    向量的每一个元素为 二进制字符串，如 '10101010'
    '''
    vec = []
    for i in range(num_in_pop):
        tmp = np.random.randint(0,256)
        # 产生一个二进制字符串。
        bin_string = bin(tmp)[2:].zfill(8)
        vec.append(bin_string)
    
    return vec

def vec2vec_num(vec):
    # 将二进制向量，转换为数值向量。
    vec_num = np.array(list(map(bin2dec, vec)))
    return vec_num

    
def generate_community(num, num_in_pop):
    '''
    产生群落，即产生（子）向量的（父）向量，其中 num 为种群的个数，即（子）向量的个数。
    '''
    community = [rand_binary_vec(num_in_pop) for i in range(num)]
    return community

def com2com_num(com):
    '''
    将二进制群落，由一个带有，二进制元素的向量 的（父）向量，
    转换成一个带有 数值矩阵（30X2） 的 向量（list）
    '''
    com_num = []
    for vec_bin in com:
        # 将 二进制 （子）向量转换为 数值向量。
        vec_num = vec2vec_num(vec_bin)
        # 将数值向量转换为 20X3 的矩阵。
        vec_num = vec_num.reshape(-1,3)
        com_num.append(vec_num)
    return com_num

def generate_random_IJ(size):
    '''
    产生两个数值 I,J。其中 I<J，且 I,J 属于 [0，size]
    '''
    tmp = range(size)
    I, J = random.sample(tmp, 2)
    if I>J:
        tmp = I
        I = J
        J = tmp
    return I, J

def swap_list(list1, I, J):
    '''
    交换向量的位置为 I，J 的元素。
    '''
    list1 = list1.copy()
    list1[I], list1[J] = list1[J], list1[I]

    return list1

def swap_str(s, I, J):
    '''
    交换字符串中位置为 I，J 的字符。
    '''
    s = list(s)
    s[I], s[J] = s[J], s[I]
    lst = ''.join(s)
    

    return lst
    
def slide_list(list1, I, J):
    '''
    将向量的 I:J 部分向前滑动，并将最后一个元素置前。
    '''
    if J-I == 1:
        list1 = list1.copy()
        tmp = list1[I:J]
        tmp = [list(tmp[-1])] + tmp[:-1]
        list1[I:J] = tmp

    return list1
    
def slide_str(s, I, J):
    '''
    将字符串的 I:J 部分向前滑动，并将最后一个元素置前。
    '''
    if J-I != 1:
        s = list(s) 
        tmp = s[I:J]
        tmp = list(tmp[-1]) + tmp[:-1]
        s[I:J] = tmp
        s = ''.join(s)

    return s
    

def flip_list(list1,I,J):
    '''
    将向量的 I:J 段翻转。
    '''
    list1 = list1.copy()
    list1[I:J] = list(reversed(list1[I:J]))
    
    return list1
    
def flip_str(s, I, J):
    '''
    将字符串的 I:J 段翻转。
    '''
    s = list(s)
    s = flip_list(s, I, J)
    lst = ''.join(s)
    
    return lst

def fitness(model, x, rows, cols): 
    '''
    适应度函数，其中 model 为 cnn 模型。
    '''
    # 要应用 cnn 模型，必须转换为（n，r，c，channel）的格式。
    x = x.reshape((1,rows,cols,1))
    y = float(model.predict(x))
    return y
    
def best_select(model, com, rows, cols):
    '''
    根据适应度函数，选择最佳模型
    com 为群落
    model 为能量预测模型
    '''
    num = len(com)
    # 初始化一个临时空列表。
    arr_result = np.zeros(num)
    # 将子群落（父向量）转换为元素是：数值、矩阵形式，的向量。
    com_num = com2com_num(com)
    for i in range(num):
        # 求解各个种群的适应度（能量）
        result = fitness(model, com_num[i], rows, cols)
        # 将适应度“填”进临时变量中
        arr_result[i] = result
        
    # 找到能量最小（适应度最高）的那个种群的索引。
    idx_best = np.argmin(arr_result)
    # 最佳种群的能量。
    best_result = np.min(arr_result)
    # 最佳种群。
    best_pop = com[idx_best]

            
    return best_pop, best_result
        
def randnormalnum(exp, var, max_num=60):
    '''
    根据正态分布，产生一个正整数
    exp 为均值
    var 为方差。
    max_num：产生的正整数不得超过 max_num
    '''
    # 根据正态分布，产生数 x（x 是随机事件）
    x = random.normalvariate(exp, var)
    # 将 x 绝对值化、整数化（四舍五入） 
    x = np.round(np.abs(x))
    # 若 x ＞ max_num，则让 x 为 max_num
    if x > max_num:
        return max_num      
    return x.astype(np.int32)

def individual_mutation(pop, rate=0.1, encoding_len=8):
    '''
    根据最佳种群 pop，使用个体基因突变的处理法，产生仨个新种群。
    rate 为变异率，决定了种群中，发生基因突变的个体数。
    '''
    num_in_pop = len(pop)    # 种群中的个体数，即60
    var = np.round(num_in_pop*rate)    # 用于产生突变个体数的正态分布的方差
    exp = 0   # 用于产生突变个体数的正态分布的均值
    
    #参与变异的个体数，是一个根据正态分布产生的随机数
    n = randnormalnum(exp, var, num_in_pop)   
    
    # 从种群 pop 中，随机选出 n 个个体
    rand_idx = random.sample(range(num_in_pop),n)
    individuals = [pop[idx] for idx in rand_idx]
        
    I, J = generate_random_IJ(encoding_len)
    
    # 临时变量
    individuals_swap = []
    individuals_slide = []
    individuals_flip = []
    
    pop_swap = pop.copy()
    pop_slide = pop.copy()
    pop_flip = pop.copy()
    
    
    if n > 0:
        for i in range(n):
            # 对个体进行处理
            individual = individuals[i]
            individual_swap = swap_str(individual, I, J)   
            individual_slide = slide_str(individual, I, J) 
            individual_flip = flip_str(individual, I, J)
            
            # 用新个体代替旧个体。
            pop_swap[rand_idx[i]] =  individual_swap
            pop_slide[rand_idx[i]] = individual_slide
            pop_flip[rand_idx[i]] = individual_flip
        
    return pop_swap, pop_slide, pop_flip
        
def mix_pop(com, num_in_pop):
    '''
    com 是新群落。
    num_in_pop 是种群的个体数。
    '''
    # 产生一个临时变量
    pop = []
    
    # 混合新子群落中的所有种群
    num_com = len(com)
    for i in range(num_in_pop//3):
        # 以原子的坐标为单位（即元素要 3 个 3 个地换）
        # 从 [0，6] 中选一个随机数
        chosen_idx = np.random.randint(0, num_com)
        # 选择随机数对应的种群
        chosen_pop = com[chosen_idx]
        
        # 将个体的坐标“放进”临时种群上，构成一个新种群    
        pop += chosen_pop[3*i:3*(i+1)]

    return pop
    
def plot_history(history_best_fitness, best_result_history, best_at_epoch,
                epochs, best_pop_history_dec):
    '''最后的画图代码'''
    fig = plt.figure(figsize=(12,6))

    plt.subplots_adjust(left=0.125, bottom=None, right=0.95, top=None,
                    wspace=0.1, hspace=None)
    
    # 画优化过程图
    plt.subplot(131)
    plt.plot(range(epochs), history_best_fitness, 'r',linewidth=1.6)
    plt.xlabel('epoch',fontsize=16)
    plt.plot(best_at_epoch, best_result_history ,'b*',markersize=20)
    plt.ylabel('best_energy', fontsize=16)
    plt.title('optimizing process',fontsize=16)
    
    # 画出最优结构的散点图
    ax1 = fig.add_subplot(132, projection='3d')
    ax1.scatter(xs = best_pop_history_dec[:,0], 
                ys=best_pop_history_dec[:,1], 
                zs=best_pop_history_dec[:,2])
    plt.title('best au25 structure', fontsize=16)
    
    # 画最优结构的“surface”图。
    ax2 = fig.add_subplot(133, projection='3d')
    ax2.plot_trisurf(best_pop_history_dec[:,0],
                    best_pop_history_dec[:,1],
                    best_pop_history_dec[:,2]\
                    , linewidth=0.2, antialiased=True)
    plt.title('best au25 structure', fontsize=16)
    plt.show()
        
if __name__ == '__main__':
    model = load_model(r'./model_and_data/cnn_model.h5')
    pop_num = 80    # 种群个数
    rows = 20
    cols = 3
    num_in_pop = rows*cols    #60个
    epochs = 500    # 遗传算法的迭代总数
    
    if pop_num%8 != 0:
        raise Exception('种群数必须能被8整除')
    
    # 构建一个记录局部最佳适应度的临时变量    
    history_best_fitness = []
    # 构建一个全局最高种群的变量
    best_pop_history = None
    # 构造一个记录最佳种群，对应的 epoch 的变量
    best_at_epoch = 0
    # 构造一个全局最高适应度的变量
    best_result_history = 9999999
    
    # 初始化群落，群落是一个 带有 80 个长度为 3X20=60 的 list 的 parent list
    # 下面将 3X20=60 的 list 称为种群，即 population
    community = generate_community(pop_num, num_in_pop)  
    
    for epoch in range(epochs):
        # 打乱群落中，种群的排列次序
        random.shuffle(community)
        
        # 子代群落
        new_community = []
        # 用于 子群落 构建的临时变量
        new_com_of8 = []
        
        # 由于下文将群落拆分成 10 个子群落，对应 10 个最佳种群，故需要记录这 10 个最佳种群
        best_pop_sub = []
        best_result_sub = []
        #print(f'群落的长度为：{len(community)}')
        #print(f'种群的个体数为：{len(community[0])}')
        
        for p in np.arange(0, pop_num, 8):
            # 以 8 个种群为单位，构建子群落（80/8=10，故10个子群落）
            tmp_com = community[p:p+8]
            # 根据适应度，从子群落中筛选最佳种群，记为 best_pop
            # print(f'子群落的长度为：{len(tmp_com)}')
            # print(f'子群落种群中的长度为{len(tmp_com[0])}')
            # print(f'所在步数{epoch}')
            
            best_pop, best_result = best_select(model, tmp_com, rows, cols)
            # print(best_pop)
            
            # 记录最佳
            best_pop_sub.append(best_pop.copy())
            best_result_sub.append(best_result)
            

            # 根据最佳种群，产生 8 个子种群，其余种群，淘汰！
            # 原封不动保留最佳种群
            new_com_of8.append(best_pop.copy())
            
            I, J = generate_random_IJ(num_in_pop)
            # 通过交换种群中的元素，即 20X3 的 list 中的元素，构建新种群。
            # 后文将称种群中的元素为：个体 
            new_pop = swap_list(best_pop.copy(), I, J)
            new_com_of8.append(new_pop)
            
            # 通过一部分个体位置“滑动”获得新种群
            new_pop = slide_list(best_pop, I, J)
            new_com_of8.append(new_pop)
            
            # 通过一部分个体位置“翻转”，获得新种群
            #print(f'个体位置翻转之前的的长度为{len(best_pop)}')
            new_pop = flip_list(best_pop.copy(), I, J)
            new_com_of8.append(best_pop)
            
            # 下面将通过修改个体的基因编码，获得新个体，构建新种群（个体位置不变）
            # 通过基因编码的片段交换，滑动，翻转获得 3 个新种群
            # rate 为变异系数，决定种群中，变异的个体数目
            pop_swap, pop_slide, pop_flip = individual_mutation(
                                    best_pop.copy(), rate=0.1, encoding_len=8
                                    )
            
            new_com_of8.append(pop_swap)
            new_com_of8.append(pop_slide)
            new_com_of8.append(pop_flip)
            
            # 通过混合上述 7 个种群，得到新种群
            pop_mix = mix_pop(new_com_of8, num_in_pop)
            
            new_com_of8.append(pop_mix)
            # 构建新群落
            new_community += new_com_of8 
            # 释放临时变量
            new_com_of8 = []
        
        # 将子代群落设置为父代
        community = new_community.copy()
        # print(f'新父代群落的长度为：{len(community)}')
        # print(f'新父代种群的个体数为：{len(community[0])}')     
        
        # 释放子代群落   
        new_community = []
        # print(community)

        # 记录全局最佳种群
        # print('全局最佳群落：',best_pop_sub)
        best_result_sub_one = min(best_result_sub)
        
        if best_result_sub_one < best_result_history:
            best_result_history = best_result_sub_one
            best_at_epoch = epoch
            min_idx = best_result_sub.index(best_result_sub_one)
            best_pop_history = best_pop_sub[min_idx]
        
        # 释放临时的，10个子群落中，10个“最佳”种群
        best_pop_sub = []
        best_result_sub = []
        history_best_fitness.append(best_result_sub_one)
        
    print('最佳适应度（最小能量）：',best_result_history)
    
    # 加载标准化模型（见2021 Mathorcup B题能量预测模型构造博文）
    scaler = pickle.load(open(r'./model_and_data/scaler.pkl', 'rb'))
    # 反标准化
    best_result_history_after = scaler.inverse_transform(
                                np.array(best_result_history).reshape(1,)
                                )
    print('最佳适应度（最小能量），标准化后：',best_result_history_after)
    
    # 将最优种群由 元素为二进制形式的向量的向量，转换为一个，数值矩阵的向量
    best_pop_history_dec = vec2vec_num(best_pop_history)
    best_pop_history_dec = best_pop_history_dec.reshape(-1,3)
    
    print('最佳结构各个点的坐标为：', best_pop_history_dec)
    
    # 将最优种群的数值矩阵向量形式保存
    pickle.dump(best_pop_history_dec, 
                open(r'./model_and_data/best_au20_structure.pkl','wb'))
        
    plot_history(history_best_fitness, best_result_history, best_at_epoch,
                epochs,best_pop_history_dec)

